import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import os
from data_preprocessing import load_data

# 设置全局字体为无衬线字体，更好地支持各种环境
plt.rcParams['font.sans-serif'] = ['Arial', 'DejaVu Sans', 'Liberation Sans', 'Bitstream Vera Sans', 'sans-serif']

def create_output_dir(dir_path='../static/visualizations'):
    """创建输出目录"""
    if not os.path.exists(dir_path):
        os.makedirs(dir_path)
    return dir_path

def plot_sensor_distributions(df, output_dir):
    """绘制各传感器数据分布图"""
    print("绘制传感器数据分布图...")
    
    # 获取传感器列
    temp_cols = [col for col in df.columns if 'Temp' in col]
    light_cols = [col for col in df.columns if 'Light' in col]
    sound_cols = [col for col in df.columns if 'Sound' in col]
    
    # 绘制温度传感器分布图
    if temp_cols:
        plt.figure(figsize=(12, 6))
        for col in temp_cols:
            sns.kdeplot(df[col], label=col)
        plt.title("Temperature Sensors Distribution")
        plt.xlabel("Temperature")
        plt.ylabel("Density")
        plt.legend()
        plt.savefig(os.path.join(output_dir, "01_温度传感器分布.png"))
        plt.close()
    
    # 绘制光线传感器分布图
    if light_cols:
        plt.figure(figsize=(12, 6))
        for col in light_cols:
            sns.kdeplot(df[col], label=col)
        plt.title("Light Sensors Distribution")
        plt.xlabel("Light Intensity")
        plt.ylabel("Density")
        plt.legend()
        plt.savefig(os.path.join(output_dir, "02_光线传感器分布.png"))
        plt.close()
    
    # 绘制声音传感器分布图
    if sound_cols:
        plt.figure(figsize=(12, 6))
        for col in sound_cols:
            sns.kdeplot(df[col], label=col)
        plt.title("Sound Sensors Distribution")
        plt.xlabel("Sound Intensity")
        plt.ylabel("Density")
        plt.legend()
        plt.savefig(os.path.join(output_dir, "03_声音传感器分布.png"))
        plt.close()

def plot_occupancy_distribution(df, output_dir):
    """绘制房间占用人数分布图"""
    print("绘制房间占用人数分布图...")
    
    plt.figure(figsize=(10, 6))
    ax = sns.countplot(x='Room_Occupancy_Count', data=df)
    
    # 添加数值标签
    for p in ax.patches:
        ax.annotate(f'{p.get_height()}', 
                    (p.get_x() + p.get_width() / 2., p.get_height()),
                    ha = 'center', va = 'bottom',
                    fontsize=10)
    
    plt.title("Room Occupancy Count Distribution")
    plt.xlabel("Number of People")
    plt.ylabel("Frequency")
    plt.savefig(os.path.join(output_dir, "04_房间占用人数分布.png"))
    plt.close()
    
    # 按小时统计平均占用人数
    if 'Hour' in df.columns:
        hourly_occupancy = df.groupby('Hour')['Room_Occupancy_Count'].mean().reset_index()
        
        plt.figure(figsize=(12, 6))
        ax = sns.barplot(x='Hour', y='Room_Occupancy_Count', data=hourly_occupancy)
        
        # 添加数值标签
        for p in ax.patches:
            ax.annotate(f'{p.get_height():.2f}', 
                        (p.get_x() + p.get_width() / 2., p.get_height()),
                        ha = 'center', va = 'bottom',
                        fontsize=8)
        
        plt.title("Average Occupancy by Hour")
        plt.xlabel("Hour of Day")
        plt.ylabel("Average Occupancy")
        plt.savefig(os.path.join(output_dir, "05_各小时段平均占用人数.png"))
        plt.close()

def plot_correlation_heatmap(df, output_dir):
    """绘制相关性热力图"""
    print("绘制相关性热力图...")
    
    # 选择数值列，但限制数量以避免图表过于复杂
    key_features = ['S1_Temp', 'S1_Light', 'S1_Sound', 'S5_CO2', 'S6_PIR', 
                     'Room_Occupancy_Count']
    
    # 如果存在新增的特征，也添加进来
    for feature in ['Temp_Mean', 'Light_Mean', 'Sound_Mean', 'PIR_Sum']:
        if feature in df.columns:
            key_features.append(feature)
    
    # 确保所有特征都存在于数据集中
    key_features = [f for f in key_features if f in df.columns]
    
    # 计算相关系数
    corr = df[key_features].corr()
    
    # 绘制热力图
    plt.figure(figsize=(12, 10))
    mask = np.triu(np.ones_like(corr, dtype=bool))  # 创建上三角掩码
    sns.heatmap(corr, annot=True, cmap='coolwarm', fmt=".2f", 
                linewidths=0.5, mask=mask)
    plt.title("Feature Correlation Heatmap")
    plt.tight_layout()
    plt.savefig(os.path.join(output_dir, "06_特征相关性热力图.png"))
    plt.close()
    
    # 与目标变量的相关性
    if 'Room_Occupancy_Count' in df.columns:
        # 获取与目标变量的相关性并排序
        target_corr = corr['Room_Occupancy_Count'].drop('Room_Occupancy_Count')
        target_corr = target_corr.abs().sort_values(ascending=False)
        
        plt.figure(figsize=(10, 6))
        target_corr.plot(kind='bar')
        plt.title("Feature Correlation with Occupancy Count")
        plt.xlabel("Features")
        plt.ylabel("Absolute Correlation")
        plt.xticks(rotation=45, ha='right')
        plt.tight_layout()
        plt.savefig(os.path.join(output_dir, "07_特征与占用人数相关性.png"))
        plt.close()

def plot_time_series(df, output_dir):
    """绘制时间序列图"""
    print("绘制时间序列图...")
    
    if 'DateTime' in df.columns:
        # 确保DateTime是日期时间类型
        if not pd.api.types.is_datetime64_any_dtype(df['DateTime']):
            df['DateTime'] = pd.to_datetime(df['DateTime'])
        
        # 按时间排序
        df_sorted = df.sort_values('DateTime')
        
        # 绘制一天内的温度变化
        if len(df_sorted) > 24:  # 确保有足够的数据
            # 取第一天的数据
            start_date = df_sorted['DateTime'].dt.date.min()
            one_day_data = df_sorted[df_sorted['DateTime'].dt.date == start_date]
            
            # 温度时间序列
            temp_cols = [col for col in df.columns if 'Temp' in col]
            if temp_cols and len(temp_cols) <= 4:  # 限制传感器数量，保持图表清晰
                plt.figure(figsize=(14, 7))
                for col in temp_cols:
                    plt.plot(one_day_data['DateTime'], one_day_data[col], label=col)
                plt.title(f"Temperature Changes During One Day ({start_date})")
                plt.xlabel("Time")
                plt.ylabel("Temperature")
                plt.legend()
                plt.xticks(rotation=45)
                plt.tight_layout()
                plt.savefig(os.path.join(output_dir, "08_一天内温度变化.png"))
                plt.close()
            
            # 房间占用人数时间序列
            plt.figure(figsize=(14, 7))
            plt.plot(one_day_data['DateTime'], one_day_data['Room_Occupancy_Count'], marker='o', linestyle='-')
            plt.title(f"Room Occupancy Changes During One Day ({start_date})")
            plt.xlabel("Time")
            plt.ylabel("Number of People")
            plt.xticks(rotation=45)
            plt.grid(True, linestyle='--', alpha=0.7)
            plt.tight_layout()
            plt.savefig(os.path.join(output_dir, "09_一天内房间占用人数变化.png"))
            plt.close()

def plot_sensor_vs_occupancy(df, output_dir):
    """绘制传感器数据与占用人数的关系图"""
    print("绘制传感器数据与占用人数的关系图...")
    
    # 选择重要的传感器特征
    key_sensors = []
    
    # 添加每种类型的第一个传感器
    for sensor_type in ['Temp', 'Light', 'Sound']:
        cols = [col for col in df.columns if sensor_type in col]
        if cols:
            key_sensors.append(cols[0])
    
    # 添加CO2和PIR传感器
    if 'S5_CO2' in df.columns:
        key_sensors.append('S5_CO2')
    if 'S6_PIR' in df.columns:
        key_sensors.append('S6_PIR')
    
    # 中文传感器名称映射
    sensor_names = {
        'S1_Temp': '温度',
        'S1_Light': '光线',
        'S1_Sound': '声音',
        'S5_CO2': '二氧化碳',
        'S6_PIR': '运动'
    }
    
    # 绘制箱线图
    for i, sensor in enumerate(key_sensors):
        plt.figure(figsize=(12, 6))
        sns.boxplot(x='Room_Occupancy_Count', y=sensor, data=df)
        plt.title(f"{sensor} vs Room Occupancy")
        plt.xlabel("Number of People")
        plt.ylabel(sensor)
        # 为图表添加网格线提高可读性
        plt.grid(True, linestyle='--', alpha=0.3)
        
        # 使用中文名称保存文件
        sensor_name = sensor_names.get(sensor, sensor)
        plt.savefig(os.path.join(output_dir, f"{10+i:02d}_{sensor_name}与占用人数关系.png"))
        plt.close()
    
    # 如果存在统计特征，也绘制它们与占用人数的关系
    if 'Sound_Mean' in df.columns:
        plt.figure(figsize=(12, 6))
        sns.boxplot(x='Room_Occupancy_Count', y='Sound_Mean', data=df)
        plt.title("Average Sound vs Room Occupancy")
        plt.xlabel("Number of People")
        plt.ylabel("Average Sound Level")
        plt.grid(True, linestyle='--', alpha=0.3)
        plt.savefig(os.path.join(output_dir, "15_平均声音与占用人数关系.png"))
        plt.close()

def create_combined_visualizations(df, output_dir):
    """创建组合可视化图表，展示多个传感器的数据"""
    print("创建组合可视化图表...")
    
    # 创建一个2x2的子图布局
    fig, axes = plt.subplots(2, 2, figsize=(16, 12))
    
    # 绘制温度与占用人数的关系
    if 'S1_Temp' in df.columns:
        sns.boxplot(x='Room_Occupancy_Count', y='S1_Temp', data=df, ax=axes[0, 0])
        axes[0, 0].set_title("Temperature vs Occupancy")
        axes[0, 0].set_xlabel("Number of People")
        axes[0, 0].set_ylabel("Temperature")
    
    # 绘制光线与占用人数的关系
    if 'S1_Light' in df.columns:
        sns.boxplot(x='Room_Occupancy_Count', y='S1_Light', data=df, ax=axes[0, 1])
        axes[0, 1].set_title("Light vs Occupancy")
        axes[0, 1].set_xlabel("Number of People")
        axes[0, 1].set_ylabel("Light Intensity")
    
    # 绘制声音与占用人数的关系
    if 'S1_Sound' in df.columns:
        sns.boxplot(x='Room_Occupancy_Count', y='S1_Sound', data=df, ax=axes[1, 0])
        axes[1, 0].set_title("Sound vs Occupancy")
        axes[1, 0].set_xlabel("Number of People")
        axes[1, 0].set_ylabel("Sound Level")
    
    # 绘制CO2与占用人数的关系
    if 'S5_CO2' in df.columns:
        sns.boxplot(x='Room_Occupancy_Count', y='S5_CO2', data=df, ax=axes[1, 1])
        axes[1, 1].set_title("CO2 vs Occupancy")
        axes[1, 1].set_xlabel("Number of People")
        axes[1, 1].set_ylabel("CO2 Level")
    
    plt.tight_layout()
    plt.savefig(os.path.join(output_dir, "16_多传感器与占用人数关系组合图.png"))
    plt.close()

def create_visualizations(data_path=None):
    """创建所有可视化图表"""
    print("开始创建数据可视化...")
    
    # 加载数据
    if data_path:
        df = pd.read_csv(data_path)
    else:
        # 尝试先加载处理后的数据，如果不存在则加载原始数据
        try:
            df = pd.read_csv('../data/first_clean/processed_room_occupancy.csv')
            print("已加载处理后的数据")
        except FileNotFoundError:
            try:
                df = pd.read_csv('../data/processed_room_occupancy.csv')
                print("已加载处理后的数据")
            except FileNotFoundError:
                df = load_data()
                print("已加载原始数据")
    
    # 创建输出目录
    output_dir = create_output_dir()
    
    # 检查是否有DateTime列，如果没有且有Date和Time列，则创建
    if 'DateTime' not in df.columns and 'Date' in df.columns and 'Time' in df.columns:
        df['DateTime'] = pd.to_datetime(df['Date'] + ' ' + df['Time'], format='%d-%m-%Y %H:%M:%S', dayfirst=True)
    
    # 生成各种可视化图表
    plot_sensor_distributions(df, output_dir)
    plot_occupancy_distribution(df, output_dir)
    plot_correlation_heatmap(df, output_dir)
    plot_time_series(df, output_dir)
    plot_sensor_vs_occupancy(df, output_dir)
    create_combined_visualizations(df, output_dir)
    
    print(f"可视化图表已保存至: {output_dir}")
    print("请打开这些图表以了解数据的特点和关系")

if __name__ == "__main__":
    create_visualizations() 